#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/5/24 16:36
@desc: 
"""
from pathlib import Path

import h5py
import pandas as pd

from utils.config import DataField as DF
from utils.config import Files


class NoduleListWriter(object):

    def __init__(self, nodule_list, file_root=Files.NODULE_ROOT):
        self.nodule_list = nodule_list
        self.file_root = Path(file_root)
        if not self.file_root.exists():
            self.file_root.mkdir()
        self.nodule_csv_file = self.file_root / Files.NODULES
        self.feature_csv_file = self.file_root / Files.FEATURE

    def _write_nodules2csv(self):
        df = pd.DataFrame()
        df[DF.SERIES_ID] = [nodule.series_id for nodule in self.nodule_list]
        df[DF.NODULE_ID] = [i for i in range(len(self.nodule_list))]
        df[DF.PATIENT_ID] = [nodule.patient_id for nodule in self.nodule_list]
        df[DF.DIAMETER] = [nodule.diameter for nodule in self.nodule_list]
        df[DF.CENTROID] = [nodule.centroid for nodule in self.nodule_list]
        file = Path(self.nodule_csv_file)
        df.to_csv(str(self.nodule_csv_file), mode='a',
                  header=not file.exists(), index=False)

    def _write_feature(self):

        def __df(index, nodule):
            sub_df = pd.concat(
                [pd.DataFrame(value, columns=[key]) for key, value in
                 nodule.features.items()], axis=1)
            sub_df[DF.PATIENT_ID] = nodule.patient_id
            sub_df[DF.SERIES_ID] = nodule.series_id
            sub_df[DF.NODULE_ID] = index
            return sub_df

        df = pd.concat(
            [__df(*values) for values in enumerate(self.nodule_list)],
            ignore_index=True)
        file = Path(self.feature_csv_file)
        df.to_csv(str(self.feature_csv_file), mode='a',
                  header=not file.exists(), index=False)

    def _write_nodules2h5py(self):
        for idx, nodule in enumerate(self.nodule_list):
            file_name = str(
                self.file_root / (nodule.series_id + '_' + str(idx) + '.h5'))
            with h5py.File(file_name, 'w') as file:
                file.create_dataset(DF.POINTS, data=nodule.points)
                file.create_dataset(DF.PATIENT_ID, data=nodule.patient_id)
                file.create_dataset(DF.SERIES_ID, data=nodule.series_id)
                file.create_dataset(DF.CENTROID, data=nodule.centroid)
                file.create_dataset(DF.DIAMETER, data=nodule.diameter)

    def update(self):
        self._write_feature()
        self._write_nodules2csv()
        self._write_nodules2h5py()
        return
